import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms
import os
import onnxruntime as ort
import numpy as np
from PIL import Image
from io import BytesIO

def remove_background(image_stream, output_stream):
    # 设置图像大小和预处理方式
    image_size = (1024, 1024)
    transform_image = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 将文件流转换为图像
    image = Image.open(image_stream)

    # 将图像转换为 RGB（如果是 RGBA 图像则会去掉 alpha 通道）
    image = image.convert("RGB")
    
    # 预处理图像
    input_image = transform_image(image).unsqueeze(0).numpy()


    # 获取项目根目录路径
    base_dir = os.path.dirname(os.path.abspath(__file__))  # 当前文件的绝对路径
    onnx_path = os.path.join(base_dir,'./photo_model.onnx')  # 将相对路径转换为绝对路径
    # 创建 ONNX Runtime 推理会话，指定使用 CUDA GPU
    session = ort.InferenceSession(onnx_path)

    # 获取模型输入的名称
    input_name = session.get_inputs()[0].name

    # 进行推理
    outputs = session.run(None, {input_name: input_image})

    # 获取模型输出并进行后处理
    output = outputs[0]  # 获取预测结果
    output = np.squeeze(output)  # 去掉多余的维度

    # 将输出转换为图像
    pred_pil = Image.fromarray((output * 255).astype(np.uint8))  # 假设输出是归一化的概率图
    mask = pred_pil.resize(image.size)
    image.putalpha(mask)

    # 将处理后的图像保存到输出文件流
    image.save(output_stream, format='PNG')

